import os
import sys
import torch
import argparse
import numpy as np
import matplotlib.pyplot as plt
import cv2
import random
import torch.utils.data as data
import torchvision.transforms as transforms
# from loss import grad_loss, ints_loss
from loss import loss_total
from MyDataLoader import VISIRData
from MyOption import args

from model import SDNet as net

#----------------------------------------------------------------
import torch.nn as nn





'''
# model2_1
    para={
    'coeff_edge': 0.05,           
    'coeff_fus_swin': 1.2, #初始16
    'coeff_vgg': 0.2, #初始1.6
    'coeff_decom': 4.5, #初始0.5
    }'''
'''
para={
    'coeff_edge': 0.05,           
    'coeff_fus_swin': 1.2, #初始16
    'coeff_vgg': 0.2, #初始1.6
    'coeff_decom': 5.5, #初始0.5
    }
'''


para={
    'coeff_edge': 0.05,           
    'coeff_fus_swin': 2.5, #初始16
    'coeff_vgg': 0.5, #初始1.6
    'coeff_decom': 5.5, #初始0.5
    'coeff_mi': 1.0    }

gpu=0
if torch.cuda.is_available():
        torch.cuda.set_device(gpu)
        device = 'cuda'
else:
        print('WARNING: [CUDA unavailable] Using CPU instead!')
        device = 'cpu'
from kornia.losses import SSIMLoss
from kornia.filters import SpatialGradient
from kornia.color import rgb_to_grayscale,ycbcr_to_rgb, rgb_to_ycbcr,grayscale_to_rgb
from kornia.enhance import equalize, invert
from SwinFusion.models.loss_vif import fusion_loss_vif,fusion_loss_vifv1
from vgg_loss import vgg_loss_fun,get_vgg_model
from trans_loss import dcom_loss
from torch.optim import lr_scheduler

MSELoss = nn.MSELoss()  
L1Loss = nn.L1Loss()
Loss_ssim = SSIMLoss(11, reduction='mean')
#criteria_fusion = Fusionloss()
swinLoss=fusion_loss_vif()
swinLossv1=fusion_loss_vifv1()
vgg_model=get_vgg_model(device=device).eval()
 
#print(vgg_model)
def loss_fun(ref_vis,ref_ir,ref_edge,int_out,edge_out):
    #---输入处理
    gray_vis=ref_vis
    y_vis=ref_vis
    #int_out_rgb=ycbcr_to_rgb(torch.cat((int_out,ycbcr_vis[:,1:,:,:]),dim=1))

    joint_in_rvEdge = ((1 - int_out)* ref_edge)
    joint_in_noEdge = int_out * (1 - ref_edge)
    #fus_edge=joint_in_noEdge+joint_in_rvEdge
    fus_edge=torch.clamp(joint_in_noEdge+joint_in_rvEdge,0,1 )

    #ycbcr_vis=rgb_to_ycbcr(ref_vis)
    #int_out_rgb=ycbcr_to_rgb(torch.cat((fus_edge,ycbcr_vis[:,1:,:,:]),dim=1))
    #---计算损失


    edge_loss=L1Loss(ref_edge,edge_out)+MSELoss(ref_edge,edge_out)
    swin_loss,_,_,_=swinLoss(gray_vis,ref_ir,fus_edge)
    vgg_loss=vgg_loss_fun(grayscale_to_rgb(equalize(gray_vis)),grayscale_to_rgb(int_out),vgg_model)
    decom_loss=dcom_loss(fus_edge,gray_vis)+dcom_loss(fus_edge,invert(ref_ir))
    '''print('decom_loss',decom_loss)
    print('vgg_loss',vgg_loss)
    print('swin_loss',swin_loss)
    print('edge_loss',edge_loss)'''
    loss =+para['coeff_edge']*edge_loss\
          +para['coeff_fus_swin']*swin_loss\
          +para['coeff_decom']*decom_loss\
          +para['coeff_vgg']*vgg_loss
          
    

    return loss

#----------------------------------------------------------------
# warnings.filterwarnings("ignore")


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True


def main():
    # 设置随机数种子
    setup_seed(args.seed)
    model_path = args.model_save_path + '/' + str(args.epoch) + '/'
    os.makedirs(model_path, exist_ok=True)
    os.makedirs(args.img_save_dir, exist_ok=True)
    # os.makedirs('./modelsave')

    lr = args.lr
    lr_patience=5
    lr_factor=0.3

    # device handling
    if args.DEVICE == 'cpu':
        device = 'cpu'
    else:
        device = args.DEVICE
        # os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    # prepare model folder
    os.makedirs(args.temp_dir, exist_ok=True)

    transform = transforms.Compose([transforms.ToTensor()])

    train_set = VISIRData(set_prefix='/home/ubuntu/workspace/data/msrs/',split='train')
    train_loader = data.DataLoader(train_set,
                                   batch_size=args.batch_size,
                                   shuffle=True,
                                   drop_last=True,
                                   num_workers=4,
                                   pin_memory=True)
    # model = net(img_size=args.imgsize, dim=256)
    model = net()

    # prepare the model for training and send to device
    model.to(device)
    model.train()
    opt = torch.optim.AdamW(model.parameters(), lr)
    scheduler = lr_scheduler.ReduceLROnPlateau(opt, mode='min', patience=lr_patience,factor=lr_factor)

    loss_plt = []
    for epoch in range(0, args.epoch):
        # os.makedirs(args.result + '/' + '%d' % (epoch + 1), exist_ok=True)

        loss_mean = []
        for idx, datas in enumerate(train_loader):
            # print(len(data))
            img1, img2, refedge = datas
            # 训练模型
            model, img_fusion,edge, loss_per_img = train(model, opt, img1, img2, refedge,lr, device)
            loss_mean.append(loss_per_img)

        # print loss
        sum_list = 0
        for item in loss_mean:
            sum_list += item
        sum_per_epoch = sum_list / len(loss_mean)
        print('Epoch--%d\tLoss:%.5f' % (epoch + 1, sum_per_epoch))
        loss_plt.append(sum_per_epoch.detach().cpu().numpy())

        scheduler.step(sum_per_epoch)
        
        # save info to txt file
        strain_path = args.temp_dir + '/temp_loss.txt'
        Loss_file = 'Epoch--' + str(epoch + 1) + '\t' + 'Loss:' + str(sum_per_epoch.detach().cpu().numpy())
        with open(strain_path, 'a') as f:
            f.write(Loss_file + '\r\n')

        if (epoch + 1) % 5 == 0:
            # 保存模型
            torch.save(model.state_dict(), model_path + args.model_save_name)
            print('model save in %s' % model_path)

    # 保存模型
    torch.save(model.state_dict(), model_path + args.model_save_name)
    print('model save in %s' % model_path)

    # 输出损失函数曲线
    plt.figure()
    x = range(0, args.epoch)  # x和y的维度要一样
    y = loss_plt
    plt.plot(x, y, 'r-')  # 设置输出样式
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.savefig(model_path + '/loss.png')  # 保存训练损失曲线图片
    plt.show()  # 显示曲线


def train(model, opt,img1, img2, refedge,lr, device):
    model.to(device)
    model.train()

    img1 = img1.to(device)
    img2 = img2.to(device)
    refedge=refedge.to(device)

    

    img1 = (img1 - img1.min()) / (img1.max() - img1.min())
    img2 = (img2 - img2.min()) / (img2.max() - img2.min())

    img_fusion, x1_de, x2_de,edge = model(img1, img2)
    # img_fusion = img_fusion.cpu()
    edge=edge.to(device)
    img_fusion = img_fusion.to(device)
    x1_de = x1_de.to(device)
    x2_de = x2_de.to(device)

    # loss_total = compute_loss(img_fusion, img_cat, img_s, img_f)
    #loss_ori = loss_total(img_fusion, x1_de, x2_de, img1, img2)
    loss_3=loss_fun(img1,img2,refedge,img_fusion,edge)
    #print('loss_ori',loss_ori)
    #print('loss_3',loss_3)
    #loss=1*loss_ori+0.1*loss_3
    #loss=0.5*loss_oril
    loss=loss_3
    opt.zero_grad()
    loss.backward()
    opt.step()

    return model, img_fusion,edge, loss


if __name__ == '__main__':
    main()
